// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/future/cancel.h>

#include <atomic>
#include <mutex>
#include <sstream>
#include <thread>
#include <utility>


#include <turbo/bootstrap/atfork.h>
#include <turbo/log/logging.h>
#include <nebula/future/mutex.h>
#include <turbo/base/macros.h>
#include <turbo/utility/signal.h>
#include <turbo/files/io_util.h>

namespace nebula {

#if ATOMIC_INT_LOCK_FREE != 2
#error Lock-free atomic int required for signal safety
#endif

    // NOTE: We care mainly about the making the common case (not cancelled) fast.

    struct StopSourceImpl {
        std::atomic<int> requested_{0};  // will be -1 or signal number if requested
        std::mutex mutex_;
        turbo::Status cancel_error_;
    };

    StopSource::StopSource() : impl_(new StopSourceImpl) {}

    StopSource::~StopSource() = default;

    void StopSource::RequestStop() { RequestStop(turbo::cancelled_error("Operation cancelled")); }

    void StopSource::RequestStop(turbo::Status st) {
        std::lock_guard<std::mutex> lock(impl_->mutex_);
                DKCHECK(!st.ok());
        if (!impl_->requested_) {
            impl_->requested_ = -1;
            impl_->cancel_error_ = std::move(st);
        }
    }

    void StopSource::RequestStopFromSignal(int signum) {
        // Only async-signal-safe code allowed here
        impl_->requested_.store(signum);
    }

    void StopSource::reset() {
        std::lock_guard<std::mutex> lock(impl_->mutex_);
        impl_->cancel_error_ = turbo::OkStatus();
        impl_->requested_.store(0);
    }

    StopToken StopSource::token() { return StopToken(impl_); }

    bool StopToken::is_stop_requested() const {
        if (!impl_) {
            return false;
        }
        return impl_->requested_.load() != 0;
    }

    turbo::Status StopToken::Poll() const {
        if (!impl_) {
            return turbo::OkStatus();
        }
        if (!impl_->requested_.load()) {
            return turbo::OkStatus();
        }

        std::lock_guard<std::mutex> lock(impl_->mutex_);
        if (impl_->cancel_error_.ok()) {
            auto signum = impl_->requested_.load();
                    DKCHECK_GT(signum, 0);
            impl_->cancel_error_ = turbo::cancelled_with_signal_payload(signum, "Operation cancelled");
        }
        return impl_->cancel_error_;
    }

    namespace {

        struct SignalStopState : public std::enable_shared_from_this<SignalStopState> {
            struct SavedSignalHandler {
                int signum;
                turbo::SignalHandler handler;
            };

            // NOTE: shared_from_this() doesn't work from constructor
            void init() {
                // XXX this pattern appears in several places, factor it out?
                atfork_handler_ = std::make_shared<turbo::AtForkHandler>(
                        /*before=*/
                        [weak_self = std::weak_ptr<SignalStopState>(shared_from_this())] {
                            auto self = weak_self.lock();
                            if (self) {
                                self->BeforeFork();
                            }
                            return self;
                        },
                        /*parent_after=*/
                        [](std::any token) {
                            auto self = std::any_cast<std::shared_ptr<SignalStopState>>(std::move(token));
                            self->ParentAfterFork();
                        },
                        /*child_after=*/
                        [](std::any token) {
                            auto self = std::any_cast<std::shared_ptr<SignalStopState>>(std::move(token));
                            self->ChildAfterFork();
                        });
                turbo::register_at_fork(atfork_handler_);
            }

            turbo::Status RegisterHandlers(const std::vector<int> &signals) {
                std::lock_guard<std::mutex> lock(mutex_);
                if (!saved_handlers_.empty()) {
                    return turbo::invalid_argument_error("Signal handlers already registered");
                }
                if (!self_pipe_) {
                    // Make sure the self-pipe is initialized
                    // (NOTE: avoid std::atomic_is_lock_free() which may require libatomic)
#if ATOMIC_POINTER_LOCK_FREE != 2
                    return turbo::unimplemented_error(
                        "Cannot setup signal StopSource because atomic pointers are not "
                        "lock-free on this platform");
#else
                    TURBO_MOVE_OR_RAISE(self_pipe_, turbo::SelfPipe::create(/*signal_safe=*/true));
#endif
                }
                if (!signal_receiving_thread_) {
                    // Spawn thread for receiving signals
                    SpawnSignalReceivingThread();
                }
                self_pipe_ptr_.store(self_pipe_.get());
                for (int signum: signals) {
                    TURBO_MOVE_OR_RAISE(auto handler,
                                        set_signal_handler(signum, turbo::SignalHandler{&HandleSignal}));
                    saved_handlers_.push_back({signum, handler});
                }
                return turbo::OkStatus();
            }

            void UnregisterHandlers() {
                std::lock_guard<std::mutex> lock(mutex_);
                self_pipe_ptr_.store(nullptr);
                auto handlers = std::move(saved_handlers_);
                for (const auto &h: handlers) {
                    KCHECK_OK(set_signal_handler(h.signum, h.handler).status());
                }
            }

            ~SignalStopState() {
                atfork_handler_.reset();
                UnregisterHandlers();
                Disable();
                if (signal_receiving_thread_) {
                    // tell the receiving thread to stop
                    auto st = self_pipe_->shutdown();
                    TURBO_WARN_NOT_OK(st, "Failed to shutdown self-pipe");
                    if (st.ok()) {
                        signal_receiving_thread_->join();
                    } else {
                        signal_receiving_thread_->detach();
                    }
                }
            }

            StopSource *stop_source() {
                std::lock_guard<std::mutex> lock(mutex_);
                return stop_source_.get();
            }

            bool enabled() {
                std::lock_guard<std::mutex> lock(mutex_);
                return stop_source_ != nullptr;
            }

            void Enable() {
                std::lock_guard<std::mutex> lock(mutex_);
                stop_source_ = std::make_shared<StopSource>();
            }

            void Disable() {
                std::lock_guard<std::mutex> lock(mutex_);
                stop_source_.reset();
            }

            static SignalStopState *instance() {
                static std::shared_ptr<SignalStopState> instance = []() {
                    auto ptr = std::make_shared<SignalStopState>();
                    ptr->init();
                    return ptr;
                }();
                return instance.get();
            }

        private:
            void SpawnSignalReceivingThread() {
                signal_receiving_thread_ = std::make_unique<std::thread>(ReceiveSignals, self_pipe_);
            }

            static void HandleSignal(int signum) {
                auto self = instance();
                if (self) {
                    self->DoHandleSignal(signum);
                }
            }

            void DoHandleSignal(int signum) {
                // async-signal-safe code only
                turbo::SelfPipe *self_pipe = self_pipe_ptr_.load();
                if (self_pipe) {
                    self_pipe->send(/*payload=*/signum);
                }
                turbo::reinstate_signal_handler(signum, &HandleSignal);
            }

            static void ReceiveSignals(std::shared_ptr<turbo::SelfPipe> self_pipe) {
                // Wait for signals on the self-pipe and propagate them to the current StopSource
                        DKCHECK(self_pipe);
                while (true) {
                    auto maybe_payload = self_pipe->wait();
                    if (turbo::is_invalid_argument(maybe_payload.status())) {
                        // Pipe shut down
                        return;
                    }
                    if (!maybe_payload.ok()) {
                        maybe_payload.status().warn();
                        return;
                    }
                    const int signum = static_cast<int>(maybe_payload.value_unsafe());
                    instance()->ReceiveSignal(signum);
                }
            }

            void ReceiveSignal(int signum) {
                std::lock_guard<std::mutex> lock(mutex_);
                if (stop_source_) {
                    stop_source_->RequestStopFromSignal(signum);
                }
            }

            // At-fork handlers

            void BeforeFork() { mutex_.lock(); }

            void ParentAfterFork() { mutex_.unlock(); }

            void ChildAfterFork() {
                new(&mutex_) std::mutex;
                // Leak previous thread, as it has become invalid.
                // We can't spawn a new one here as it would have unfortunate side effects;
                // especially in the frequent context of a fork+exec.
                // (for example the Python subprocess module closes all fds before calling exec)
                TURBO_UNUSED(signal_receiving_thread_.release());
                // Make internal state consistent: with no listening thread, we shouldn't
                // feed the self-pipe from the signal handler.
                UnregisterHandlers();
            }

            std::mutex mutex_;
            std::vector<SavedSignalHandler> saved_handlers_;
            std::shared_ptr<StopSource> stop_source_;
            std::unique_ptr<std::thread> signal_receiving_thread_;
            std::shared_ptr<turbo::AtForkHandler> atfork_handler_;

            // For signal handler interaction
            std::shared_ptr<turbo::SelfPipe> self_pipe_;
            // Raw atomic pointer, as atomic load/store of a shared_ptr may not be lock-free
            // (it is not on libstdc++).
            std::atomic<turbo::SelfPipe *> self_pipe_ptr_;
        };

    }  // namespace

    turbo::Result<StopSource *> SetSignalStopSource() {
        auto stop_state = SignalStopState::instance();
        if (stop_state->enabled()) {
            return turbo::invalid_argument_error("Signal stop source already set up");
        }
        stop_state->Enable();
        return stop_state->stop_source();
    }

    void ResetSignalStopSource() {
        auto stop_state = SignalStopState::instance();
                DKCHECK(stop_state->enabled());
        stop_state->Disable();
    }

    turbo::Status RegisterCancellingSignalHandler(const std::vector<int> &signals) {
        auto stop_state = SignalStopState::instance();
        if (!stop_state->enabled()) {
            return turbo::invalid_argument_error("Signal stop source was not set up");
        }
        return stop_state->RegisterHandlers(signals);
    }

    void UnregisterCancellingSignalHandler() {
        auto stop_state = SignalStopState::instance();
                DKCHECK(stop_state->enabled());
        stop_state->UnregisterHandlers();
    }

}  // namespace nebula
